library(Seurat)
library(tidyverse)
library(ggpubr)
library(pheatmap)
library(tidyseurat)
library(patchwork)

housekeep <- c('ACTB', 'GAPDH', 'RPL13A', 'RPS18', 'EEF1A1', 'RPLP0', 'UBC', 'YWHAZ')

crc10x <- read_rds('CRC-I/data/Zhang-Yu-2020/tpm_10x.rds')
crc_smart <- read_rds('CRC-I/data/Zhang-Yu-2020/tpm_smart.rds')

# housekeeping genes in crc ------
get_expr_pos <- function(sobj, features, name){
  hk_expr <- FetchData(sobj, features)
  
  prcnt_list <- map_dbl(hk_expr, \(x)PercentAbove(x, 0)) %>%
    formattable::percent() %>%
    as.character()
  
  prcnt_list <- tibble(pos = prcnt_list, gene = names(prcnt_list))
  
  hk_expr %>%
    pivot_longer(everything(), names_to = 'gene', values_to = 'log(tpm+1)') %>%
    left_join(prcnt_list) %>%
    add_column(platform = name)}

expr_10x <- get_expr_pos(crc10x, housekeep, '10X')
expr_smart <- get_expr_pos(crc_smart, housekeep, 'Smart-seq2')

expr_crc <- expr_10x %>%
  bind_rows(expr_smart) 

write_rds(expr_crc, 'DE_cells/results/vln_pos_crc.rds')
expr_crc <- read_rds('DE_cells/results/vln_pos_crc.rds')

expr_crc %>%
  ggplot(aes(gene, `log(tpm+1)`, fill = platform)) +
  geom_violin() +
  geom_text(aes(label = pos, y = 3.2), data = distinct(expr_crc, pos, .keep_all = TRUE), size = 5) +
  expand_limits(y = 4) +
  facet_wrap(~platform, scales = 'free')+
  theme_pubr() +
  labs_pubr()

crc_housekeep <- last_plot()

# B cell markers --------
b_cell_marker <- read_table('DE_cells/ref/B_cell_marker.txt', col_names = FALSE) %>% pull(X1)

split_and_pos <- function(sobj, features, name) {
  bcm_10x_list <- sobj %>%
    subset(Global_Cluster == 'B cell') %>%
    SplitObject('Sub_Cluster') %>%
    map(\(x)get_expr_pos(x, features, name))
  
  bcm_10x <- bcm_10x_list %>%
    map2(names(bcm_10x_list), \(x, y)add_column(x, cell_type = y)) %>%
    reduce(bind_rows)
}

bcm_10x <- split_and_pos(crc10x, b_cell_marker, '10X')
bcm_smart <- split_and_pos(crc_smart, b_cell_marker, 'Smart-seq2')

cell_name <- tibble(
  cell_type = unique(crc_cd19$cell_type),
  B_cell_subset = c('IgD+ follicular B cells',
                    'MS4A1+ follicular B cells',
                    'LRMP+ GC B cells',
                    'IgG+ Plasma cells',
                    'IgA+ GALT B cells')
)

crc_cd19 <- bcm_10x %>%
  bind_rows(bcm_smart) %>%
  filter(gene == 'CD19') %>%
  left_join(cell_name)

write_rds(crc_cd19, 'DE_cells/results/crc_cd19.rds')
crc_cd19 <- read_rds('DE_cells/results/crc_cd19.rds')

crc_cd19 %>%
  ggplot(aes(str_wrap(B_cell_subset, width = 12), `log(tpm+1)`, fill = platform)) +
  geom_violin() +
  geom_text(aes(label = pos, y = 3.2), data = distinct(crc_cd19, pos, .keep_all = TRUE), size = 5) +
  expand_limits(y = 4) +
  facet_wrap(~platform, scales = 'free')+
  theme_pubr() +
  labs_pubr() +
  labs(x = 'B cell subset', title = 'CD19')

cd19_vln_crc <- last_plot()

# housekeeping & B cell marker in HC as heatmap ---------
blish_sobj <- read_rds('DE_cells/data/blish_cleaned.rds')
bgi_sobj <- read_rds('DE_cells/data/BGI_cleaned.rds')

blish_sobj <- subset(blish_sobj, Status == "Healthy" & scrublet_call == 'Singlet')
bgi_sobj <- subset(bgi_sobj, Stage == "Ctrl" & scrublet_call == 'Singlet')

cross_tibble <- expand_grid(
  sobj = c('blish', 'bgi'),
  feature = list(housekeep, b_cell_marker))

find_hook <- function(hook){
  switch (hook,
          'bgi' = bgi_sobj,
          'blish' = blish_sobj)
}

plot_hooked_heatmap <- function(sobj, features){
  pos_percent_for_types <- function(type_name){
    find_hook(sobj) %>%
      subset(ident = type_name) %>%
      FetchData(vars = features) %>%
      map(\(x)PercentAbove(x, 0)) %>%
      as.data.frame() %>%
      add_column(cell_type = type_name)
  }
  
  find_hook(sobj) %>%
    Idents() %>%
    unique() %>%
    lapply(pos_percent_for_types) %>%
    purrr::reduce(bind_rows) ->
    heat_mat
  
  heat_mat %>%
    column_to_rownames('cell_type') %>%
    mutate(across(everything(), ~formatC(.*100, format = 'f', digits = 2))) %>%
    mutate(across(everything(), ~paste0(.,'%'))) ->
    percent_mat_chr
  
  heatmat_fig <- heat_mat %>%
    column_to_rownames('cell_type') %>%
    t() %>%
    pheatmap(
      cluster_row = FALSE,
      cluster_cols = FALSE,
      display_numbers = t(percent_mat_chr),
      number_color = "black",
      #legend = FALSE,
      fontsize = 8,
      fontsize_number = case_when(
        sobj == 'bgi' ~ 8,
        sobj == 'blish' ~ 7
      )
    ) %>%
    ggplotify::as.ggplot()
  
  list(heatmat_fig, percent_mat_chr)}

map2(cross_tibble$sobj,
     cross_tibble$feature,
     plot_hooked_heatmap) %>%
  transpose() ->
  heatmap_list

heatmap_source <- heatmap_list[[2]]

heatmap_list <- heatmap_list[[1]]

# final fig5: housekeep -------
crc_housekeep /
  (heatmap_list[[3]] + heatmap_list[[1]]) +
  plot_annotation(tag_levels = 'A') +
  plot_layout(heights = c(1,2)) &
  theme(plot.tag = element_text(size = 32))

ggsave("DE_cells/figures/fig5_ms_housekeep.pdf",
       width = 36,
       height = 30,
       units = 'cm')

ggsave("DE_cells/figures/fig5_ms_housekeep.png",
       width = 36,
       height = 30,
       units = 'cm')

# final fig6: B cell markers -------
cd19_vln_crc /
  (heatmap_list[[4]] + heatmap_list[[2]]) +
  patchwork::plot_annotation(tag_levels = 'A') +
  plot_layout(heights = c(1,2.5)) &
  theme(plot.tag = element_text(size = 32))

ggsave("DE_cells/figures/fig6_ms_cd19_B_markers.pdf",
       width = 36,
       height = 30,
       units = 'cm')

ggsave("DE_cells/figures/fig6_ms_cd19_B_markers.png",
       width = 36,
       height = 30,
       units = 'cm')

# save source data ------
source_path <- mutate(cross_tibble, path = str_glue('DE_cells/results/fig5-6_source_{sobj}_{feature}.csv')) %>%
  pull(path)

walk2(heatmap_source, source_path, write.csv)
